import sys
import os
import torch
import numpy as np
import random

sys.path.append(os.path.join(os.path.dirname(__file__), "../src"))
from mlapo_torch import MLAPO
from mlapo_triton import mlapo_triton
from mlapo_gen_cmp_results import gen_test_input, TEST_CASES, TEST_CASE_HEADER

sys.path.append(os.path.join(os.path.dirname(__file__), "../../../precision"))
from compare import check_operator_accuracy

sys.path.append(os.path.join(os.path.dirname(__file__), "../../../profiler"))
from triton_profiler import triton_profiler_wrapper, times_list_header

sys.path.append(os.path.join(os.path.dirname(__file__), "../../../utils"))
# from utils import write_csv
from templates import prof_template


# 设置运行卡号
torch.npu.set_device(0)  
# 随机种子
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)


def prof_mlapo(
    test_case
):
    # get seperated input
    input_param = gen_test_input(*test_case)

    # profiling
    def call(): 

        # call torch and triton
        # mlapo = MLAPO()
        # golden_output = mlapo.mla_preprocess_calc(*input_param)

        actual_output = mlapo_triton(*input_param)
    
    # de bench
    times_list = triton_profiler_wrapper(call)
    print(times_list)
    return times_list


if __name__ == "__main__":
    
    csv_header_list = TEST_CASE_HEADER + times_list_header
    prof_template("Mlapo", TEST_CASES, prof_mlapo, csv_header_list)


    
    

    